#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: zyx
@Date: 2024/11/29 15:19
@FileName: app.py
@Description: web页面
"""

from qa_service import ChemQA
import gradio as gr
import argparse
from utils import get_session

def parse():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--rag-dir',
        required=True,
        help='rag文件所在目录',
        type=str,
    )
    parser.add_argument(
        '-m',
        '--memory-type',
        choices=['memory', 'redis'],
        default='memory',
        help='langchain历史记录存储方式 memory:内存 redis:redis存储'
    )
    parser.add_argument(
        '-r',
        '--rag-type',
        choices=['vector_query', 'lightrag'],
        default='lightrag',
        help='langchain rag 方式 vector_query:使用chroma进行相似度计算 lightrag:使用lightrag知识图谱构建'
    )
    # 解析参数
    return parser.parse_args()

class ChemApp:

    def __init__(self) -> None:
        args = parse()
        print(args)
        self.qa = ChemQA(rag_dir=args.rag_dir, rag_type=args.rag_type, memory_type=args.memory_type)
        css = """
        .gradio-container { max-width:850px !important; margin:10px auto !important;}
        .message { padding: 5px !important; font-size: 14px !important;}
        """
        title = """
        <center> 
        <h3> 化学AI机器人 </h3>
        </center>
        """
        self.app = gr.Blocks(css=css)
        with self.app:
            session_id = gr.State(get_session)
            with gr.Row():
                gr.HTML(title)
            chatbot = gr.Chatbot(type="messages")
            msg = gr.Textbox(autofocus=True)
            with gr.Row():
                submitBtn = gr.Button("提交问题", variant="primary")
                submitBtn.click(self.chat, [msg, chatbot, session_id], [msg, chatbot])
                gr.ClearButton([msg, chatbot, session_id], value="清除记录")
            gr.Examples(
                examples=["如何使用AI进行产率预测", "分子设计与发现有哪些常用方法"],
                inputs=msg,
                outputs=msg,
            )

    def chat(self, message, chat_history, session_id):
        bot_response = self.qa.answer(message, session_id)
        # gradio 历史展示通过 user 及 assistant 两个角色
        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": bot_response})
        # response返回的第一个空值，是为了将输入msg内容置零
        return "", chat_history

    def start(self):
        self.app.launch(share=False)


if __name__ == "__main__":
    app = ChemApp()
    app.start()
